"""
简单的使用示例：不使用 Agent，直接使用 LLM
支持切换 Qwen 和 PanGu 模型
"""
from qwen_llm import QwenLLM
from pangu_llm import PanguLLM
from typing import Literal, Optional


def simple_chat_example(
    model_type: Optional[Literal["qwen", "pangu"]] = None,
    model_path: Optional[str] = None
):
    """简单的对话示例
    
    Args:
        model_type: 模型类型，"qwen" 或 "pangu"，如果不指定则让用户选择
        model_path: 可选的模型路径（覆盖默认路径）
    """
    print("=" * 50)
    print("简单对话示例 - 支持 Qwen 和 PanGu 模型")
    print("=" * 50)
    
    
    # 如果未指定模型类型，让用户选择
    if model_type is None:
        print("\n请选择要使用的模型:")
        print("1. Qwen3-Omni (默认)")
        print("2. PanGu-Alpha")
        
        choice = input("\n请输入选择 (1/2, 默认为1): ").strip()
        
        if choice == "2":
            model_type = "pangu"
            model_name = "PanGu-Alpha"
        else:
            model_type = "qwen"
            model_name = "Qwen3-Omni"
    else:
        model_name = "PanGu-Alpha" if model_type == "pangu" else "Qwen3-Omni"
    
    print(f"\n使用模型: {model_name}")
    print("\n正在加载模型...")
    
    # 根据模型类型初始化 LLM
    if model_type == "pangu":
        if model_path:
            llm = PanguLLM(model_path=model_path)
        else:
            llm = PanguLLM()
    else:  # 默认使用 Qwen
        if model_path:
            llm = QwenLLM(model_path=model_path)
        else:
            llm = QwenLLM()
    
    print("\n模型已加载！开始对话（输入 'quit' 退出）\n")
    
    while True:
        try:
            user_input = input("你: ")
            
            if user_input.lower() in ['quit', 'exit', '退出']:
                print("再见!")
                break
            
            if not user_input.strip():
                continue
            
            # 调用 LLM
            print("思考中...")
            response = llm.invoke(user_input)
            print(f"\n助手: {response}\n")
            
        except KeyboardInterrupt:
            print("\n\n程序中断。再见!")
            break
        except Exception as e:
            print(f"\n错误: {str(e)}\n")


if __name__ == "__main__":
    simple_chat_example()

